from paddlex import create_model
import numpy as np
import sys
import os
# 获取到 config.py 所在的目录（相对于 get_table_cell.py 的路径）
current_dir = os.path.dirname(os.path.abspath(__file__))
config_dir = os.path.abspath(os.path.join(current_dir, '../../..'))  # 返回到 src 目录

# 将 config.py 所在的路径添加到 sys.path
sys.path.append(config_dir)
from config import config
LABEL_NAME = config.label_name
THRESHOLD = config.threshold

def get_table_label(img) -> str:
    """
    根据图片路径预测表格的类型，返回表格标签（'wired_table' 或 'wireless_table'）。
    :param img: Python Var/str/list,
        Python变量，如numpy.ndarray表示的图像数据
        文件路径，如图像文件的本地路径：/root/data/img.jpg
        URL链接，如图像文件的网络URL：示例
        本地目录，该目录下需包含待预测数据文件，如本地路径：/root/data/
        列表，列表元素需为上述类型数据，如[numpy.ndarray, numpy.ndarray]，[\"/root/data/img1.jpg\", \"/root/data/img2.jpg\"]，[\"/root/data1\", \"/root/data2\"]
    """
    label_model = create_model(model_name="PP-LCNet_x1_0_table_cls")
    label_output = label_model.predict(img, batch_size=1)
    label_output = list(label_output)[0]
    
    # 获取分数最高的标签
    max_score_index = np.argmax(label_output['scores'])
    label_name = label_output['label_names'][max_score_index]
    
    return label_name



def get_cell_coordinates(img, label_name: str = LABEL_NAME, threshold: float = THRESHOLD) -> list:
    """
    根据表格标签获取表格单元格的坐标列表。
    
    :param img: Python Var/str/list,
        Python变量，如numpy.ndarray表示的图像数据
        文件路径，如图像文件的本地路径：/root/data/img.jpg
        URL链接，如图像文件的网络URL：示例
        本地目录，该目录下需包含待预测数据文件，如本地路径：/root/data/
        列表，列表元素需为上述类型数据，如[numpy.ndarray, numpy.ndarray]，[\"/root/data/img1.jpg\", \"/root/data/img2.jpg\"]，[\"/root/data1\", \"/root/data2\"]
    :param label_name: 表格标签 ('wired_table' 或 'wireless_table')
    :param threshold: 用于预测时过滤的置信度阈值，默认为0.3
    :return: 包含单元格坐标的列表
    """
    if label_name not in ['wired_table', 'wireless_table','auto']:
        raise ValueError(f"Invalid label_name '{label_name}'. Expected 'wired_table' 、'auto' or 'wireless_table'.")
    
    if label_name == 'auto':
        label_name = get_table_label(img)
    if label_name == 'wired_table':
        model = create_model(model_name="RT-DETR-L_wired_table_cell_det")
    else:
        model = create_model(model_name="RT-DETR-L_wireless_table_cell_det")
    
    # 获取单元格坐标
    cell_coords = model.predict(img, threshold=threshold, batch_size=1)
    cell_boxes = list(cell_coords)[0]['boxes']
    
    # 提取坐标
    cell_coordinates = [boxe['coordinate'] for boxe in cell_boxes]
    
    return cell_coordinates

